{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "My2AZpV_RuTs"
      },
      "source": [
        "# Multi-modal\n",
        "\n",
        "[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google-deepmind/gemma/blob/main/colabs/multimodal.ipynb)\n",
        "\n",
        "Example on how to use Gemma models for multi-modal."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "94CVV9ZxKVDO"
      },
      "outputs": [],
      "source": [
        "!pip install -q gemma"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PXBc1hRKRuTt"
      },
      "outputs": [],
      "source": [
        "# Common imports\n",
        "import os\n",
        "import jax.numpy as jnp\n",
        "import tensorflow_datasets as tfds\n",
        "\n",
        "# Gemma imports\n",
        "from gemma import gm"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "i_DEVehe3v5Y"
      },
      "source": [
        "By default, Jax do not utilize the full GPU memory, but this can be overwritten. See [GPU memory allocation](https://docs.jax.dev/en/latest/gpu_memory_allocation.html):"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AaK17GWo3v5Y"
      },
      "outputs": [],
      "source": [
        "os.environ[\"XLA_PYTHON_CLIENT_MEM_FRACTION\"]=\"1.00\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wN3Ve867t_o7"
      },
      "source": [
        "First, let's load an image:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "height": 541
        },
        "executionInfo": {
          "elapsed": 1381,
          "status": "ok",
          "timestamp": 1741415862715,
          "user": {
            "displayName": "",
            "userId": ""
          },
          "user_tz": -60
        },
        "id": "-H4yUpuDmAla",
        "outputId": "716ca3f7-30c0-4dd3-cb0b-f506a483f3b1"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "\u003cstyle\u003e\n",
              "      .ndarray_repr .ndarray_raw_data {\n",
              "        display: none;\n",
              "      }\n",
              "      .ndarray_repr.show_array .ndarray_raw_data {\n",
              "        display: block;\n",
              "      }\n",
              "      .ndarray_repr.show_array .ndarray_image_preview {\n",
              "        display: none;\n",
              "      }\n",
              "      \u003c/style\u003e\n",
              "      \u003cdiv id=\"id-acdf0a9a-a317-4fee-8fb2-1913bfcd158a\" class=\"ndarray_repr\"\u003e\u003cpre\u003endarray (500, 667, 3) \u003cbutton style=\"padding: 0 2px;\"\u003eshow data\u003c/button\u003e\u003c/pre\u003e\u003cimg src=\"\" class=\"ndarray_image_preview\" /\u003e\u003cpre class=\"ndarray_raw_data\"\u003earray([[[1, 1, 0],\n",
              "        [1, 1, 0],\n",
              "        [1, 1, 0],\n",
              "        ...,\n",
              "        [1, 1, 1],\n",
              "        [1, 1, 1],\n",
              "        [1, 1, 1]],\n",
              "\n",
              "       [[1, 1, 0],\n",
              "        [1, 1, 0],\n",
              "        [1, 1, 0],\n",
              "        ...,\n",
              "        [1, 1, 1],\n",
              "        [1, 1, 1],\n",
              "        [1, 1, 1]],\n",
              "\n",
              "       [[1, 1, 0],\n",
              "        [1, 1, 0],\n",
              "        [1, 1, 0],\n",
              "        ...,\n",
              "        [1, 1, 1],\n",
              "        [1, 1, 1],\n",
              "        [1, 1, 1]],\n",
              "\n",
              "       ...,\n",
              "\n",
              "       [[1, 1, 1],\n",
              "        [1, 1, 1],\n",
              "        [1, 1, 1],\n",
              "        ...,\n",
              "        [7, 7, 5],\n",
              "        [6, 6, 4],\n",
              "        [5, 5, 3]],\n",
              "\n",
              "       [[1, 1, 1],\n",
              "        [1, 1, 1],\n",
              "        [1, 1, 1],\n",
              "        ...,\n",
              "        [7, 7, 5],\n",
              "        [6, 6, 4],\n",
              "        [5, 5, 3]],\n",
              "\n",
              "       [[1, 1, 1],\n",
              "        [1, 1, 1],\n",
              "        [1, 1, 1],\n",
              "        ...,\n",
              "        [7, 7, 5],\n",
              "        [6, 6, 4],\n",
              "        [5, 5, 3]]], shape=(500, 667, 3), dtype=uint8)\u003c/pre\u003e\u003c/div\u003e\u003cscript\u003e\n",
              "      (() =\u003e {\n",
              "      const titles = ['show data', 'hide data'];\n",
              "      let index = 0\n",
              "      document.querySelector('#id-acdf0a9a-a317-4fee-8fb2-1913bfcd158a button').onclick = (e) =\u003e {\n",
              "        document.querySelector('#id-acdf0a9a-a317-4fee-8fb2-1913bfcd158a').classList.toggle('show_array');\n",
              "        index = (++index) % 2;\n",
              "        document.querySelector('#id-acdf0a9a-a317-4fee-8fb2-1913bfcd158a button').textContent = titles[index];\n",
              "        e.preventDefault();\n",
              "        e.stopPropagation();\n",
              "      }\n",
              "      })();\n",
              "    \u003c/script\u003e"
            ],
            "text/plain": [
              "array([[[1, 1, 0],\n",
              "        [1, 1, 0],\n",
              "        [1, 1, 0],\n",
              "        ...,\n",
              "        [1, 1, 1],\n",
              "        [1, 1, 1],\n",
              "        [1, 1, 1]],\n",
              "\n",
              "       [[1, 1, 0],\n",
              "        [1, 1, 0],\n",
              "        [1, 1, 0],\n",
              "        ...,\n",
              "        [1, 1, 1],\n",
              "        [1, 1, 1],\n",
              "        [1, 1, 1]],\n",
              "\n",
              "       [[1, 1, 0],\n",
              "        [1, 1, 0],\n",
              "        [1, 1, 0],\n",
              "        ...,\n",
              "        [1, 1, 1],\n",
              "        [1, 1, 1],\n",
              "        [1, 1, 1]],\n",
              "\n",
              "       ...,\n",
              "\n",
              "       [[1, 1, 1],\n",
              "        [1, 1, 1],\n",
              "        [1, 1, 1],\n",
              "        ...,\n",
              "        [7, 7, 5],\n",
              "        [6, 6, 4],\n",
              "        [5, 5, 3]],\n",
              "\n",
              "       [[1, 1, 1],\n",
              "        [1, 1, 1],\n",
              "        [1, 1, 1],\n",
              "        ...,\n",
              "        [7, 7, 5],\n",
              "        [6, 6, 4],\n",
              "        [5, 5, 3]],\n",
              "\n",
              "       [[1, 1, 1],\n",
              "        [1, 1, 1],\n",
              "        [1, 1, 1],\n",
              "        ...,\n",
              "        [7, 7, 5],\n",
              "        [6, 6, 4],\n",
              "        [5, 5, 3]]], shape=(500, 667, 3), dtype=uint8)"
            ]
          },
          "execution_count": 5,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "ds = tfds.data_source('oxford_flowers102', split='train')\n",
        "image = ds[0]['image']\n",
        "image"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XzEj3PYvX1Sm"
      },
      "source": [
        "Load the model and params."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "u_YP1u-dx4sx"
      },
      "outputs": [],
      "source": [
        "model = gm.nn.Gemma3_4B()\n",
        "\n",
        "params = gm.ckpts.load_params(gm.ckpts.CheckpointPath.GEMMA3_4B_IT)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AH0eWFWJaiNk"
      },
      "source": [
        "## Sampling full prompt\n",
        "\n",
        "To use the multi-modal capabilities, simply:\n",
        "\n",
        "* In the prompt: Add the `\u003cstart_of_image\u003e` special tokens, where the images should be inserted.\n",
        "* Pass the image(s) to the `images=` argument of the sampler"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "executionInfo": {
          "elapsed": 8012,
          "status": "ok",
          "timestamp": 1741214002398,
          "user": {
            "displayName": "",
            "userId": ""
          },
          "user_tz": -60
        },
        "id": "6R5J42EiZtkC",
        "outputId": "ffc9047b-9c8b-42d1-c092-8f922e9bc745"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Here's a breakdown of what I can say about the image:\n",
            "\n",
            "**Overall Impression:**\n",
            "\n",
            "The image is a stunning, close-up photograph of a water lily in full bloom. It’s dramatically lit, creating a strong contrast between light and shadow, which really emphasizes the flower's form and texture.\n",
            "\n",
            "**Specific Details:**\n",
            "\n",
            "*   **Flower Type:** It appears to be a Nymphaea (water lily). The shape of the petals and the prominent stamens are characteristic of this type of flower.\n",
            "*   **Color:** The petals are primarily white with a subtle pinkish hue at the base. The stamens are a bright, vibrant yellow.\n",
            "*   **Lighting:** The lighting is key. There's a strong light source coming from the upper left, casting dramatic shadows and highlighting the edges of the petals. This creates a sense of depth and makes the flower appear almost sculptural.\n",
            "*   **Texture:** You can see the delicate texture of the petals – they appear smooth but with subtle ridges and folds.\n",
            "*   **Composition:** The flower is centered in the frame, drawing the viewer's eye directly to it. The dark background isolates the flower and makes it the focal point.\n",
            "*   **Water Droplets:** There are a few water droplets on the petals, adding a touch of freshness and realism.\n",
            "\n",
            "**Mood/Feeling:**\n",
            "\n",
            "The image evokes a feeling of tranquility, beauty, and perhaps a touch of mystery due to the dramatic lighting. It feels serene and peaceful.\n",
            "\n",
            "**Do you want me to focus on a specific aspect of the image, such as:**\n",
            "\n",
            "*   The lighting technique?\n",
            "*   The flower's anatomy?\n",
            "*   The overall mood it creates?\u003cend_of_turn\u003e\n"
          ]
        }
      ],
      "source": [
        "sampler = gm.text.ChatSampler(\n",
        "    model=model,\n",
        "    params=params,\n",
        ")\n",
        "\n",
        "out = sampler.chat(\n",
        "    'What can you say about this image: \u003cstart_of_image\u003e',\n",
        "    images=image,\n",
        ")\n",
        "print(out)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ykeLumksxdCT"
      },
      "source": [
        "Notes:\n",
        "\n",
        "* The model was trained on `jpeg` images. If you have PNG images, those should be encoded/decoded to Jpeg, to avoid bias.\n",
        "\n",
        "* You can pass multiple images. Just add `\u003cstart_of_image\u003e` everywhere an image should be inserted. All images should be resized to the same shape. Input shape would then be `batch, num_images, h, w, c` (instead of `batch, h, w, c`).\n",
        "\n",
        "* If prompts within a batch have different number of images, just pad the tensor with 0 (or any) values for unused images."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TPf01271RuTs"
      },
      "source": [
        "## Calling model directly\n",
        "\n",
        "Adding images to the model only require to:\n",
        "\n",
        "* In the prompt: Add `\u003cstart_of_image\u003e` special tokens where the images should be inserted."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "executionInfo": {
          "elapsed": 222,
          "status": "ok",
          "timestamp": 1741417624790,
          "user": {
            "displayName": "",
            "userId": ""
          },
          "user_tz": -60
        },
        "id": "P0_bPukcdca5"
      },
      "outputs": [],
      "source": [
        "tokenizer = gm.text.Gemma3Tokenizer()\n",
        "\n",
        "\n",
        "prompt = \"\"\"\u003cstart_of_turn\u003euser\n",
        "Describe this image in a single word.\n",
        "\n",
        "\u003cstart_of_image\u003e\n",
        "\n",
        "\u003cend_of_turn\u003e\n",
        "\u003cstart_of_turn\u003emodel\n",
        "\"\"\"\n",
        "prompt = jnp.asarray(tokenizer.encode(prompt, add_bos=True))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EdDyedvik1Af"
      },
      "source": [
        "* In the model: pass the `images=` to `model.apply`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "height": 597
        },
        "executionInfo": {
          "elapsed": 13393,
          "status": "ok",
          "timestamp": 1741213887254,
          "user": {
            "displayName": "",
            "userId": ""
          },
          "user_tz": -60
        },
        "id": "G3Kbo9hgRuTt",
        "outputId": "8581ce89-4d13-4a37-ab02-13e2a2d2f713"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "\u003chtml\u003e\n",
              "\u003chead\u003e\u003cmeta charset=\"utf-8\" /\u003e\u003c/head\u003e\n",
              "\u003cbody\u003e\n",
              "    \u003cdiv\u003e            \u003cscript src=\"https://cdnjs.cloudflare.com/ajax/libs/mathjax/2.7.5/MathJax.js?config=TeX-AMS-MML_SVG\"\u003e\u003c/script\u003e\u003cscript type=\"text/javascript\"\u003eif (window.MathJax \u0026\u0026 window.MathJax.Hub \u0026\u0026 window.MathJax.Hub.Config) {window.MathJax.Hub.Config({SVG: {font: \"STIX-Web\"}});}\u003c/script\u003e                \u003cscript type=\"text/javascript\"\u003ewindow.PlotlyConfig = {MathJaxConfig: 'local'};\u003c/script\u003e\n",
              "        \u003cscript charset=\"utf-8\" src=\"https://cdn.plot.ly/plotly-2.34.0.min.js\"\u003e\u003c/script\u003e                \u003cdiv id=\"f6426dad-2356-4f37-b142-a88cf1788b04\" class=\"plotly-graph-div\" style=\"height:525px; width:100%;\"\u003e\u003c/div\u003e            \u003cscript type=\"text/javascript\"\u003e                                    window.PLOTLYENV=window.PLOTLYENV || {};                                    if (document.getElementById(\"f6426dad-2356-4f37-b142-a88cf1788b04\")) {                    Plotly.newPlot(                        \"f6426dad-2356-4f37-b142-a88cf1788b04\",                        [{\"alignmentgroup\":\"True\",\"hovertemplate\":\"x=%{x}\\u003cbr\\u003ey=%{y}\\u003cextra\\u003e\\u003c\\u002fextra\\u003e\",\"legendgroup\":\"\",\"marker\":{\"color\":\"#636efa\",\"pattern\":{\"shape\":\"\"}},\"name\":\"\",\"offsetgroup\":\"\",\"orientation\":\"v\",\"showlegend\":false,\"textposition\":\"auto\",\"x\":[\"'Lotus'\",\"'Elegant'\",\"'Bloom'\",\"'Ser'\",\"'Magn'\",\"'Floral'\",\"'Frag'\",\"'Beauty'\",\"'Bl'\",\"'L'\",\"'Lily'\",\"'E'\",\"'Pet'\",\"'Maj'\",\"'Grace'\",\"'Del'\",\"'Ref'\",\"'Blooming'\",\"'Flower'\",\"'Lil'\",\"'Op'\",\"'P'\",\"'Radi'\",\"'Beautiful'\",\"'Or'\",\"'Spl'\",\"'Water'\",\"'Flora'\",\"'Blo'\",\"'Res'\"],\"xaxis\":\"x\",\"y\":[0.59375,0.2177734375,0.1318359375,0.037841796875,0.006591796875,0.00311279296875,0.00311279296875,0.00311279296875,0.0024261474609375,0.00189208984375,0.000888824462890625,0.000255584716796875,0.000255584716796875,0.0001544952392578125,0.0001201629638671875,9.393692016601562e-05,5.698204040527344e-05,4.4345855712890625e-05,4.4345855712890625e-05,3.457069396972656e-05,1.633167266845703e-05,1.2695789337158203e-05,5.990266799926758e-06,4.649162292480469e-06,4.649162292480469e-06,3.635883331298828e-06,3.635883331298828e-06,1.7210841178894043e-06,1.043081283569336e-06,1.043081283569336e-06],\"yaxis\":\"y\",\"type\":\"bar\"}],                        {\"template\":{\"data\":{\"histogram2dcontour\":[{\"type\":\"histogram2dcontour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"choropleth\":[{\"type\":\"choropleth\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"histogram2d\":[{\"type\":\"histogram2d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmap\":[{\"type\":\"heatmap\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmapgl\":[{\"type\":\"heatmapgl\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"contourcarpet\":[{\"type\":\"contourcarpet\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"contour\":[{\"type\":\"contour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"surface\":[{\"type\":\"surface\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"mesh3d\":[{\"type\":\"mesh3d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"scatter\":[{\"fillpattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2},\"type\":\"scatter\"}],\"parcoords\":[{\"type\":\"parcoords\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolargl\":[{\"type\":\"scatterpolargl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"bar\":[{\"error_x\":{\"color\":\"#2a3f5f\"},\"error_y\":{\"color\":\"#2a3f5f\"},\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"bar\"}],\"scattergeo\":[{\"type\":\"scattergeo\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolar\":[{\"type\":\"scatterpolar\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"histogram\":[{\"marker\":{\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"histogram\"}],\"scattergl\":[{\"type\":\"scattergl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatter3d\":[{\"type\":\"scatter3d\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattermapbox\":[{\"type\":\"scattermapbox\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterternary\":[{\"type\":\"scatterternary\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattercarpet\":[{\"type\":\"scattercarpet\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"carpet\":[{\"aaxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"baxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"type\":\"carpet\"}],\"table\":[{\"cells\":{\"fill\":{\"color\":\"#EBF0F8\"},\"line\":{\"color\":\"white\"}},\"header\":{\"fill\":{\"color\":\"#C8D4E3\"},\"line\":{\"color\":\"white\"}},\"type\":\"table\"}],\"barpolar\":[{\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"barpolar\"}],\"pie\":[{\"automargin\":true,\"type\":\"pie\"}]},\"layout\":{\"autotypenumbers\":\"strict\",\"colorway\":[\"#636efa\",\"#EF553B\",\"#00cc96\",\"#ab63fa\",\"#FFA15A\",\"#19d3f3\",\"#FF6692\",\"#B6E880\",\"#FF97FF\",\"#FECB52\"],\"font\":{\"color\":\"#2a3f5f\"},\"hovermode\":\"closest\",\"hoverlabel\":{\"align\":\"left\"},\"paper_bgcolor\":\"white\",\"plot_bgcolor\":\"#E5ECF6\",\"polar\":{\"bgcolor\":\"#E5ECF6\",\"angularaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"radialaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"ternary\":{\"bgcolor\":\"#E5ECF6\",\"aaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"baxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"caxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"coloraxis\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"colorscale\":{\"sequential\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"sequentialminus\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"diverging\":[[0,\"#8e0152\"],[0.1,\"#c51b7d\"],[0.2,\"#de77ae\"],[0.3,\"#f1b6da\"],[0.4,\"#fde0ef\"],[0.5,\"#f7f7f7\"],[0.6,\"#e6f5d0\"],[0.7,\"#b8e186\"],[0.8,\"#7fbc41\"],[0.9,\"#4d9221\"],[1,\"#276419\"]]},\"xaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"yaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"scene\":{\"xaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"yaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"zaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2}},\"shapedefaults\":{\"line\":{\"color\":\"#2a3f5f\"}},\"annotationdefaults\":{\"arrowcolor\":\"#2a3f5f\",\"arrowhead\":0,\"arrowwidth\":1},\"geo\":{\"bgcolor\":\"white\",\"landcolor\":\"#E5ECF6\",\"subunitcolor\":\"white\",\"showland\":true,\"showlakes\":true,\"lakecolor\":\"white\"},\"title\":{\"x\":0.05},\"mapbox\":{\"style\":\"light\"}}},\"xaxis\":{\"anchor\":\"y\",\"domain\":[0.0,1.0],\"title\":{\"text\":\"Tokens\"}},\"yaxis\":{\"anchor\":\"x\",\"domain\":[0.0,1.0],\"title\":{\"text\":\"Probability\"}},\"legend\":{\"tracegroupgap\":0},\"margin\":{\"t\":60},\"barmode\":\"relative\",\"title\":{\"text\":\"Probability Distribution of Tokens\"}},                        {\"responsive\": true}                    ).then(function(){\n",
              "                            \n",
              "var gd = document.getElementById('f6426dad-2356-4f37-b142-a88cf1788b04');\n",
              "var x = new MutationObserver(function (mutations, observer) {{\n",
              "        var display = window.getComputedStyle(gd).display;\n",
              "        if (!display || display === 'none') {{\n",
              "            console.log([gd, 'removed!']);\n",
              "            Plotly.purge(gd);\n",
              "            observer.disconnect();\n",
              "        }}\n",
              "}});\n",
              "\n",
              "// Listen for the removal of the full notebook cells\n",
              "var notebookContainer = gd.closest('#notebook-container');\n",
              "if (notebookContainer) {{\n",
              "    x.observe(notebookContainer, {childList: true});\n",
              "}}\n",
              "\n",
              "// Listen for the clearing of the current output cell\n",
              "var outputEl = gd.closest('.output');\n",
              "if (outputEl) {{\n",
              "    x.observe(outputEl, {childList: true});\n",
              "}}\n",
              "\n",
              "                        })                };                            \u003c/script\u003e        \u003c/div\u003e\n",
              "\u003c/body\u003e\n",
              "\u003c/html\u003e"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "# Run the model\n",
        "out = model.apply(\n",
        "    {'params': params},\n",
        "    tokens=prompt,\n",
        "    images=image,\n",
        "    return_last_only=True,  # Only predict the last token\n",
        ")\n",
        "\n",
        "\n",
        "# Plot the probability distribution\n",
        "tokenizer.plot_logits(out.logits)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "G_6j755AyZww"
      },
      "source": [
        "## Finetuning\n",
        "\n",
        "Finetuning with multi-modal is also simple. From the original [finetuning](https://github.com/google-deepmind/gemma/blob/main/docs/finetuning.md), changing to multi-modal only require 2 changes:\n",
        "\n",
        "* Have a dataset which also return an image (`b h w c`), or multiple images (`b n h w c`)\n",
        "* Specify the model input which field in the batch correspond to the images:\n",
        "\n",
        "  ```python\n",
        "  model = gm.nn.Gemma3_4B(\n",
        "      tokens='batch.tokens',\n",
        "      images='batch.image',\n",
        "  )\n",
        "  ```\n",
        "\n",
        "See the [`multimodal.py`](https://github.com/google-deepmind/gemma/tree/main/examples/multimodal.py) example."
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "last_runtime": {},
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
